#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
AI Bot 3.0 - Elasticsearch增强版
集成Elasticsearch高性能文档检索的智能助手
"""

import os
import sys
import json5
import urllib.parse
import logging
import hashlib
from typing import List, Dict, Any, Optional
from datetime import datetime

# 添加当前目录到Python路径，以便导入本地模块
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from qwen_agent.tools.base import BaseTool, register_tool
from qwen_agent.agents import Assistant
from qwen_agent.agents.elasticsearch_assistant import ElasticsearchAssistant
from qwen_agent.gui import WebUI
from qwen_agent.tools.simple_doc_parser import SimpleDocParser

try:
    from elasticsearch import Elasticsearch
    from elasticsearch.helpers import bulk
    ES_AVAILABLE = True
except ImportError:
    print("⚠️ elasticsearch库未安装，将使用默认检索模式")
    print("安装命令: pip install elasticsearch")
    ES_AVAILABLE = False

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)


@register_tool('my_image_gen')
class MyImageGen(BaseTool):
    """AI图像生成工具"""
    description = 'AI 绘画（图像生成）服务，输入文本描述，返回基于文本信息绘制的图像 URL。'
    parameters = [{
        'name': 'prompt',
        'type': 'string',
        'description': '期望的图像内容的详细描述',
        'required': True
    }]

    def call(self, params: str, **kwargs) -> str:
        prompt = json5.loads(params)['prompt']
        prompt = urllib.parse.quote(prompt)
        return json5.dumps(
            {'image_url': f'https://image.pollinations.ai/prompt/{prompt}'},
            ensure_ascii=False
        )


@register_tool('elasticsearch_retrieval')
class ElasticsearchRetrievalTool(BaseTool):
    """
    Elasticsearch检索工具 - 本地集成版
    替代qwen_agent默认的retrieval工具
    """
    
    description = '从Elasticsearch索引中检索与查询相关的保险文档片段'
    parameters = [{
        'name': 'query',
        'type': 'string',
        'description': '搜索查询字符串',
        'required': True
    }]
    
    def __init__(self, cfg: Optional[Dict] = None):
        super().__init__(cfg)
        self.cfg = cfg or {}
        
        # ES配置
        self.es_host = self.cfg.get('es_host', 'https://localhost')
        self.es_port = self.cfg.get('es_port', 9200)
        self.es_user = self.cfg.get('es_user', 'elastic')
        self.es_password = self.cfg.get('es_password', 'ncUg27UP5aMSzSElVXAX')
        self.index_name = self.cfg.get('index_name', 'qwen_agent_bot3_docs')
        self.chunk_size = self.cfg.get('chunk_size', 500)
        self.max_results = self.cfg.get('max_results', 5)
        
        # 初始化ES客户端
        self.es_client = None
        self.doc_parser = None
        
        if ES_AVAILABLE:
            self._connect_elasticsearch()
            self._create_index_if_not_exists()
            self.doc_parser = SimpleDocParser(cfg)
    
    def _connect_elasticsearch(self) -> bool:
        """连接到Elasticsearch"""
        try:
            self.es_client = Elasticsearch(
                f"{self.es_host}:{self.es_port}",
                basic_auth=(self.es_user, self.es_password),
                verify_certs=False,
                request_timeout=60,
                max_retries=3,
                retry_on_timeout=True
            )
            
            if self.es_client.ping():
                logger.info("✅ 成功连接到Elasticsearch")
                return True
            else:
                logger.error("❌ 无法连接到Elasticsearch")
                return False
                
        except Exception as e:
            logger.error(f"❌ ES连接失败: {str(e)}")
            return False
    
    def _create_index_if_not_exists(self):
        """创建索引（如果不存在）"""
        if not self.es_client:
            return
            
        try:
            if not self.es_client.indices.exists(index=self.index_name):
                mapping = {
                    "mappings": {
                        "properties": {
                            "file_path": {"type": "keyword"},
                            "file_name": {"type": "keyword"},
                            "chunk_id": {"type": "keyword"},
                            "content": {"type": "text"},
                            "metadata": {"type": "object"},
                            "created_time": {"type": "date"},
                            "file_hash": {"type": "keyword"}
                        }
                    },
                    "settings": {
                        "number_of_shards": 1,
                        "number_of_replicas": 0
                    }
                }
                
                self.es_client.indices.create(index=self.index_name, body=mapping)
                logger.info(f"✅ 创建索引: {self.index_name}")
            else:
                logger.info(f"📋 索引已存在: {self.index_name}")
                
        except Exception as e:
            logger.error(f"❌ 创建索引失败: {str(e)}")
    
    def _get_file_hash(self, file_path: str) -> str:
        """获取文件哈希值"""
        try:
            with open(file_path, 'rb') as f:
                content = f.read()
                return hashlib.md5(content).hexdigest()
        except Exception:
            return ""
    
    def _chunk_text(self, text: str) -> List[str]:
        """将文本分块"""
        if len(text) <= self.chunk_size:
            return [text]
        
        chunks = []
        start = 0
        
        while start < len(text):
            end = start + self.chunk_size
            
            # 尝试在句号处分割
            if end < len(text):
                for i in range(end, min(end + 100, len(text))):
                    if text[i] in '。！？\n':
                        end = i + 1
                        break
            
            chunk = text[start:end].strip()
            if chunk:
                chunks.append(chunk)
            
            start = end - 50  # 50字符重叠
            if start >= len(text):
                break
        
        return chunks
    
    def index_files(self, files: List[str]) -> bool:
        """索引文件列表"""
        if not ES_AVAILABLE or not self.es_client:
            logger.warning("ES不可用，跳过索引")
            return False
        
        logger.info(f"📚 开始索引 {len(files)} 个文件")
        
        actions = []
        indexed_count = 0
        
        for file_path in files:
            try:
                logger.info(f"📄 处理文件: {os.path.basename(file_path)}")
                
                # 解析文档
                if self.doc_parser:
                    documents = self.doc_parser.load(file_path)
                else:
                    # 简单文本读取
                    with open(file_path, 'r', encoding='utf-8') as f:
                        content = f.read()
                    documents = [type('Doc', (), {'page_content': content, 'metadata': {'source': file_path}})()]
                
                if not documents:
                    continue
                
                # 提取文本内容
                full_text = ""
                metadata = {}
                for doc in documents:
                    full_text += doc.page_content + "\n"
                    if hasattr(doc, 'metadata'):
                        metadata.update(doc.metadata)
                
                if not full_text.strip():
                    continue
                
                # 文本分块
                chunks = self._chunk_text(full_text)
                file_hash = self._get_file_hash(file_path)
                
                # 为每个块创建索引文档
                for i, chunk in enumerate(chunks):
                    action = {
                        "_index": self.index_name,
                        "_id": f"{file_hash}_{i}",
                        "_source": {
                            "file_path": file_path,
                            "file_name": os.path.basename(file_path),
                            "chunk_id": f"{os.path.basename(file_path)}_{i}",
                            "content": chunk,
                            "metadata": metadata,
                            "created_time": datetime.now().isoformat(),
                            "file_hash": file_hash
                        }
                    }
                    actions.append(action)
                
                indexed_count += 1
                logger.info(f"✅ 文件分块完成: {len(chunks)} 个块")
                
            except Exception as e:
                logger.error(f"❌ 处理文件失败 {file_path}: {str(e)}")
                continue
        
        # 批量索引
        if actions:
            try:
                success_count, failed_items = bulk(
                    self.es_client,
                    actions,
                    chunk_size=100,
                    request_timeout=60
                )
                
                self.es_client.indices.refresh(index=self.index_name)
                logger.info(f"✅ 批量索引完成: {success_count} 个文档块")
                return True
                
            except Exception as e:
                logger.error(f"❌ 批量索引失败: {str(e)}")
                return False
        
        return True
    
    def search(self, query: str, k: int = None) -> List[Dict[str, Any]]:
        """搜索文档"""
        if not ES_AVAILABLE or not self.es_client:
            return []
        
        if k is None:
            k = self.max_results
        
        try:
            search_body = {
                "query": {
                    "bool": {
                        "should": [
                            {
                                "match": {
                                    "content": {
                                        "query": query,
                                        "boost": 2.0
                                    }
                                }
                            },
                            {
                                "match_phrase": {
                                    "content": {
                                        "query": query,
                                        "boost": 3.0
                                    }
                                }
                            }
                        ],
                        "minimum_should_match": 1
                    }
                },
                "highlight": {
                    "fields": {
                        "content": {
                            "fragment_size": 150,
                            "number_of_fragments": 2
                        }
                    }
                },
                "size": k,
                "_source": ["file_name", "content", "metadata", "chunk_id"]
            }
            logger.info(f"🔍 执行搜索查询: {json.dumps(search_body, ensure_ascii=False, indent=2)}")
            response = self.es_client.search(
                index=self.index_name,
                body=search_body
            )
            
            results = []
            for hit in response['hits']['hits']:
                source = hit['_source']
                
                result = {
                    'page_content': source['content'],
                    'metadata': {
                        'source': source['file_name'],
                        'chunk_id': source['chunk_id'],
                        'score': hit['_score'],
                        **source.get('metadata', {})
                    }
                }
                
                if 'highlight' in hit:
                    result['metadata']['highlights'] = hit['highlight'].get('content', [])
                
                results.append(result)
            
            logger.info(f"🔍 搜索完成: 找到 {len(results)} 个相关文档")
            return results
            
        except Exception as e:
            logger.error(f"❌ 搜索失败: {str(e)}")
            return []
    
    def call(self, params: str, **kwargs) -> str:
        """工具调用接口"""
        try:
            if isinstance(params, str):
                params = json5.loads(params)
            
            query = params.get('query', '')
            if not query:
                return json5.dumps([], ensure_ascii=False)
            
            # 执行搜索
            results = self.search(query)
            
            # 格式化结果
            formatted_results = []
            for result in results:
                formatted_results.append({
                    'content': result['page_content'],
                    'source': result['metadata']['source'],
                    'score': result['metadata']['score'],
                    'chunk_id': result['metadata']['chunk_id']
                })
            
            return json5.dumps(formatted_results, ensure_ascii=False)
            
        except Exception as e:
            logger.error(f"❌ 工具调用失败: {str(e)}")
            return json5.dumps([], ensure_ascii=False)


class AIBot3:
    """AI Bot 3.0 - Elasticsearch增强版"""
    
    def __init__(self):
        self.bot = None
        self.es_config = self._get_default_es_config()
        
    def _get_default_es_config(self) -> Dict[str, Any]:
        """获取默认ES配置"""
        return {
            'es_host': 'https://localhost',
            'es_port': 9200,
            'es_user': 'elastic',
            'es_password': 'ncUg27UP5aMSzSElVXAX',
            'index_name': 'qwen_agent_bot3_docs',
            'verify_certs': False,
            'chunk_size': 500,
            'chunk_overlap': 50,
            'max_results': 5
        }
    
    def _get_files_list(self) -> List[str]:
        """获取docs文件夹下的所有文件"""
        file_dir = os.path.join(os.path.dirname(__file__), 'docs')
        files = []
        
        if os.path.exists(file_dir):
            for file in os.listdir(file_dir):
                file_path = os.path.join(file_dir, file)
                if os.path.isfile(file_path):
                    files.append(file_path)
        
        logger.info(f"📁 找到 {len(files)} 个文件: {[os.path.basename(f) for f in files]}")
        return files
    
     # ... 其他代码 ...

    def show_stats(self):
        """显示系统统计信息"""
        try:
            if not self.bot:
                self.init_agent_service()
            
            print("📊 AI Bot 3.0 统计信息")
            print("="*40)
            
            # 检索模式
            mode = self.bot.get_retrieval_mode() if hasattr(self.bot, 'get_retrieval_mode') else '默认'
            print(f"检索模式: {mode}")
            
            # 索引统计
            if hasattr(self.bot, 'get_index_stats'):
                stats = self.bot.get_index_stats()
                if 'error' not in stats:
                    print(f"索引名称: {stats.get('index_name', 'N/A')}")
                    print(f"文档块数: {stats.get('document_count', 0)}")
                    print(f"索引文件: {stats.get('indexed_files', 0)}")
                    print(f"存储大小: {stats.get('store_size', 0)} bytes")
                else:
                    print(f"索引状态: {stats.get('error', '未知错误')}")

            # 查看当前索引中的文档数量
            if hasattr(self.bot, 'es_client') and self.bot.es_client:
                try:
                    count_response = self.bot.es_client.count(index=self.index_name)
                    print(f"🔍 当前索引文档总数: {count_response['count']}")
                except Exception as e:
                    print(f"❌ 获取文档数量失败: {e}")

            # 文件列表
            files = self._get_files_list()
            print(f"可用文件: {len(files)}")
            for i, file_path in enumerate(files[:5], 1):
                print(f"  {i}. {os.path.basename(file_path)}")
            if len(files) > 5:
                print(f"  ... 还有 {len(files) - 5} 个文件")
            
            print("="*40)
            
        except Exception as e:
            print(f"❌ 获取统计信息失败: {str(e)}")
    def init_agent_service(self, use_elasticsearch: bool = True) -> ElasticsearchAssistant:
        """
        初始化AI助手服务
        
        Args:
            use_elasticsearch: 是否使用Elasticsearch检索
            
        Returns:
            ElasticsearchAssistant实例
        """
        logger.info("🚀 初始化AI Bot 3.0...")
        
        # LLM配置
        llm_cfg = {
            'model': 'qwen-max',
            'model_server': 'dashscope',
            'api_key': os.getenv('DASHSCOPE_API_KEY'),
            'generate_cfg': {
                'top_p': 0.8
            }
        }
        
        # 系统指令
        system_instruction = '''你是一个专业的保险知识AI助手。

你的主要能力包括：
1. 📚 **知识检索**: 从保险文档库中检索相关信息回答用户问题
2. 🎨 **图像生成**: 根据用户描述生成相关图像
3. 💻 **代码执行**: 执行Python代码进行数据分析和图像处理

工作流程：
1. 当用户询问保险相关问题时，首先使用检索工具查找相关文档
2. 基于检索到的信息提供准确、专业的回答
3. 如果用户需要图像，调用图像生成工具
4. 必要时使用代码执行工具进行数据处理

请始终：
- 优先使用检索到的文档信息
- 提供准确、专业的保险知识
- 用中文回复用户
- 保持友好、专业的语调'''
        
        # 工具列表
        tools = ['my_image_gen', 'code_interpreter','elasticsearch_retrieval']
        
        # 获取文件列表
        files = self._get_files_list()
        
        # 创建Elasticsearch增强的Assistant
        if use_elasticsearch:
            logger.info("🔍 启用Elasticsearch高性能检索模式")
            self.bot = ElasticsearchAssistant(
                function_list=tools,
                llm=llm_cfg,
                system_message=system_instruction,
                name="保险知识AI助手 3.0",
                description="基于Elasticsearch的高性能保险知识问答助手",
                files=files,
                use_elasticsearch=True,
                es_config=self.es_config
            )
        else:
            logger.info("📚 使用默认检索模式")
            # 降级到普通Assistant
            from qwen_agent.agents import Assistant
            self.bot = Assistant(
                function_list=tools,
                llm=llm_cfg,
                system_message=system_instruction,
                files=files
            )
        
        # 显示索引统计
        if hasattr(self.bot, 'get_index_stats'):
            stats = self.bot.get_index_stats()
            if 'error' not in stats:
                logger.info(f"📊 索引统计: {stats['document_count']} 个文档块, "
                          f"{stats['indexed_files']} 个文件")
        
        logger.info("✅ AI助手初始化完成")
        return self.bot
    
    def app_tui(self, use_elasticsearch: bool = True):
        """
        终端交互模式
        
        Args:
            use_elasticsearch: 是否使用Elasticsearch
        """
        try:
            # 初始化助手
            bot = self.init_agent_service(use_elasticsearch)
            
            print("🤖 AI Bot 3.0 - 终端交互模式")
            print("="*60)
            print(f"检索模式: {bot.get_retrieval_mode() if hasattr(bot, 'get_retrieval_mode') else '默认'}")
            print("输入 'quit' 退出程序")
            print("="*60)
            
            # 对话历史
            messages = []
            
            while True:
                try:
                    # 获取用户输入
                    query = input('\n🔍 请输入问题: ').strip()
                    
                    if query.lower() in ['quit', 'exit', '退出']:
                        print("👋 再见！")
                        break
                    
                    if not query:
                        print('⚠️ 问题不能为空！')
                        continue
                    
                    # 构建消息
                    messages.append({'role': 'user', 'content': query})
                    
                    print("🤔 正在思考...")
                    
                    # 运行助手并处理响应
                    response = []
                    current_index = 0
                    first_chunk = True
                    
                    for response_chunk in bot.run(messages=messages):
                        if first_chunk:
                            # 显示检索信息
                            if hasattr(bot, 'search_documents'):
                                try:
                                    retrieved_docs = bot.search_documents(query, k=3)
                                    if retrieved_docs:
                                        print("\n📚 检索到的相关文档:")
                                        for i, doc in enumerate(retrieved_docs, 1):
                                            source = doc.get('metadata', {}).get('source', '未知')
                                            score = doc.get('metadata', {}).get('score', 0)
                                            print(f"  {i}. {source} (相关度: {score:.3f})")
                                        print()
                                except Exception as e:
                                    logger.debug(f"显示检索信息失败: {e}")
                            
                            first_chunk = False
                        
                        # 流式输出响应
                        if response_chunk and response_chunk[0]['role'] == 'assistant':
                            assistant_message = response_chunk[0]
                            new_content = assistant_message.get('content', '')
                            print(new_content[current_index:], end='', flush=True)
                            current_index = len(new_content)
                        
                        response = response_chunk
                    
                    print()  # 换行
                    messages.extend(response)
                    
                except KeyboardInterrupt:
                    print("\n👋 程序被用户中断")
                    break
                except Exception as e:
                    print(f"❌ 处理请求时出错: {str(e)}")
                    logger.error(f"处理请求失败: {e}")
                    
        except Exception as e:
            print(f"❌ 启动终端模式失败: {str(e)}")
            logger.error(f"启动失败: {e}")
    
    def app_gui(self, use_elasticsearch: bool = True):
        """
        Web图形界面模式
        
        Args:
            use_elasticsearch: 是否使用Elasticsearch
        """
        try:
            print("🌐 启动 AI Bot 3.0 Web界面...")
            
            # 初始化助手
            bot = self.init_agent_service(use_elasticsearch)
            
            # 配置聊天界面
            retrieval_mode = bot.get_retrieval_mode() if hasattr(bot, 'get_retrieval_mode') else '默认'
            
            chatbot_config = {
                'prompt.suggestions': [
                    '介绍下雇主责任险的保障范围',
                    '工伤保险和雇主险有什么区别？',
                    '画一只正在学习保险知识的猫',
                    '企业团体意外险的理赔流程',
                    '财产一切险包含哪些保障？',
                    '帮我画一个保险公司的logo'
                ]
            }
            
            print(f"🔍 检索模式: {retrieval_mode}")
            print("🌐 Web界面准备就绪，正在启动服务...")
            
            # 启动Web界面
            WebUI(bot, chatbot_config=chatbot_config).run()
            
        except Exception as e:
            print(f"❌ 启动Web界面失败: {str(e)}")
            logger.error(f"Web界面启动失败: {e}")
            print("💡 请检查:")
            print("  1. 网络连接是否正常")
            print("  2. DASHSCOPE_API_KEY 环境变量是否设置")
            print("  3. Elasticsearch服务是否运行")
    
    def show_stats(self):
        """显示系统统计信息"""
        try:
            if not self.bot:
                self.init_agent_service()
            
            print("📊 AI Bot 3.0 统计信息")
            print("="*40)
            
            # 检索模式
            mode = self.bot.get_retrieval_mode() if hasattr(self.bot, 'get_retrieval_mode') else '默认'
            print(f"检索模式: {mode}")
            
            # 索引统计
            if hasattr(self.bot, 'get_index_stats'):
                stats = self.bot.get_index_stats()
                if 'error' not in stats:
                    print(f"索引名称: {stats.get('index_name', 'N/A')}")
                    print(f"文档块数: {stats.get('document_count', 0)}")
                    print(f"索引文件: {stats.get('indexed_files', 0)}")
                    print(f"存储大小: {stats.get('store_size', 0)} bytes")
                else:
                    print(f"索引状态: {stats.get('error', '未知错误')}")
            
            # 文件列表
            files = self._get_files_list()
            print(f"可用文件: {len(files)}")
            for i, file_path in enumerate(files[:5], 1):
                print(f"  {i}. {os.path.basename(file_path)}")
            if len(files) > 5:
                print(f"  ... 还有 {len(files) - 5} 个文件")
            
            print("="*40)
            
        except Exception as e:
            print(f"❌ 获取统计信息失败: {str(e)}")


def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description='AI Bot 3.0 - Elasticsearch增强版')
    parser.add_argument('--mode', choices=['gui', 'tui', 'stats'], default='gui',
                       help='运行模式: gui(Web界面), tui(终端), stats(统计信息)')
    parser.add_argument('--no-es', action='store_true',
                       help='禁用Elasticsearch，使用默认检索')
    
    args = parser.parse_args()
    
    # 创建AI Bot实例
    bot = AIBot3()
    
    try:
        if args.mode == 'gui':
            bot.app_gui(use_elasticsearch=not args.no_es)
        elif args.mode == 'tui':
            bot.app_tui(use_elasticsearch=not args.no_es)
        elif args.mode == 'stats':
            bot.show_stats()
    except KeyboardInterrupt:
        print("\n👋 程序被用户中断")
    except Exception as e:
        print(f"❌ 程序执行失败: {str(e)}")
        logger.error(f"程序执行失败: {e}")


if __name__ == '__main__':
    main()